# core/engine.py

# 导入所有组件
from ..components import storage, retrievers, postprocessors, models
from llama_index.core.query_engine import RetrieverQueryEngine


# 注意：之前你遇到的导入问题，这里我们假设已经修复
# 如果还报错，请使用 from fusion_retriever.components import ...

def build_query_engine():
    """
    Builds the complete query engine by integrating all components.
    """
    # ----------------------------------------------------------------------
    # 步骤 0: 初始化 LLM 和 Embedding 模型 (这是新增的关键步骤)
    # ----------------------------------------------------------------------
    models.initialize_models()

    print("🚀 [Engine] Starting to build the full query engine...")

    # 步骤 1: 初始化存储组件
    vector_store = storage.get_pg_vector_store()
    docstore = storage.get_mongo_docstore()

    # 步骤 2: 获取数据并构建索引
    all_nodes = storage.get_nodes_from_docstore(docstore)
    vector_index = storage.get_vector_index(vector_store)

    # 步骤 3: 创建检索器
    vector_retriever = retrievers.create_vector_retriever(vector_index)
    bm25_retriever = retrievers.create_bm25_retriever(all_nodes)
    fusion_retriever = retrievers.create_fusion_retriever(vector_retriever, bm25_retriever)

    # 步骤 4: 创建后处理器
    meta_post = postprocessors.get_metadata_replacement_postprocessor()
    rerank_post = postprocessors.get_sentence_transformer_reranker()

    # 步骤 5: 组装最终的查询引擎
    query_engine = RetrieverQueryEngine.from_args(
        retriever=fusion_retriever,
        node_postprocessors=[meta_post, rerank_post],
        streaming=True
    )

    print("✅ [Engine] Query engine built successfully!")
    return query_engine


# 创建全局查询引擎实例
query_engine_instance = build_query_engine()